import sys
import random
import pylab as plt
import numpy as np
import math
import os
import networkx as nx

from params import *
datasets = ['bkite','fsquare','gowalla']
lbls = ['Brightkite', 'Foursquare', 'Gowalla']
node_dist = [5651.0,8494.0,5663.0]
link_dist = [2041.0,1442.0,1792.0]
degs = [7.88,22.07,9.48,7.0]

def ecdf(data):
    l1 = math.log10(min(data))
    l2 = math.log10(max(data)) 
    bins = 10**np.linspace(l1,l2, 100)

    pdf1,bins1,x = plt.hist(data,
            bins=bins,cumulative=False,normed=False)
    plt.close()

    c = [pdf1[0]]
    for j in range(1,len(pdf1)):
        t = c[j-1]
        c.append(t+pdf1[j])

    total = c[-1]
    c1 = map(lambda x: 1.0*x/total, c)
    print 'total ', total
    return bins1[:-1],c1

def load_graph(file):
    s = 0.0
    k = 0
    i = 0
    g = nx.Graph()
    for line in open(file):
        i += 1
        if i % 1000000 == 0:
            print 'line ', i
        if line.startswith('#'):
            continue
        u1,u2,dist = line.split(' ')
        u1,u2 = map(int,(u1,u2))
        dist = float(dist)
        dist = max((1.0,dist))
        s += dist
        k += 1
        g.add_edge(u1,u2,weight=dist)
    avg_dist = s/k
    return g, avg_dist

def user_shape_factor(file,graph):
    tr = 0
    user_shapes = {}
    for line in open(file):
        #if random.random() >= 0.01:
        #    continue
        n1,n2,n3 = sorted(map(int,line.split()))
        l1 = graph[n1][n2]['weight']
        l2 = graph[n1][n3]['weight']
        l3 = graph[n2][n3]['weight']
        shape = float(l1+l2+l3)/3
        tr += 1
        def update_user(node):
            s,n = user_shapes.setdefault(node,(0.0,0))
            s += shape
            n += 1
            user_shapes[node] = (s,n)
        update_user(n1)
        update_user(n2)
        update_user(n3)
    shapes = [a/b for a,b in user_shapes.values()]
    avg_shape = sum(shapes)/len(shapes)

    x,c = ecdf(shapes)
    print 'read %d user shape factors  with average shape %f'%(
            len(shapes),avg_shape)
    return avg_shape,x,c

def shape_factor(file,graph):
    tr = 0
    shapes = []
    for line in open(file):
        #if random.random() >= 0.01:
        #    continue
        n1,n2,n3 = sorted(map(int,line.split()))
        l1 = graph[n1][n2]['weight']
        l2 = graph[n1][n3]['weight']
        l3 = graph[n2][n3]['weight']
        shapes.append(float(l1+l2+l3)/3)
        tr += 1
    avg_shape = sum(shapes)/len(shapes)
    x,c = ecdf(shapes)
    print 'read %d triangles with average shape %f'%(tr,avg_shape)
    return avg_shape,x,c

plots = []
for dataset in datasets:
    geo_trace_file = os.path.join(WORKDIR,'gsn','null_graphs',dataset,
            '%s_geo_null_graph_1.txt'%dataset)
    social_trace_file = os.path.join(WORKDIR, 'gsn','null_graphs',dataset,
            '%s_social_null_graph_2.txt'%dataset)
    social_triangle_file = os.path.join(WORKDIR, 'gsn','results',dataset,'%s_triangles.txt'%dataset)
    geo_triangle_file = os.path.join(WORKDIR, 'gsn','null_graphs',dataset,'%s_geo_null_1_triangles.txt'%dataset)
    print 'Dataset %s'%dataset
    trace_file = os.path.join(WORKDIR, 'gsn','traces',dataset,
            '%s_graph.txt'%dataset)
    triangle_file = os.path.join(WORKDIR, 'gsn','results',dataset,'%s_triangles.txt'%dataset)

    graph,avg_dist = load_graph(trace_file)
    print 'average link length ', avg_dist
    avg_shape,x,c = user_shape_factor(triangle_file,graph)
    print 'Data shape ', avg_shape/avg_dist

    del graph

    geo_graph, geo_avg_dist = load_graph(geo_trace_file)
    print "average geo link length", geo_avg_dist
    avg_geo_shape,x_geo,c_geo = user_shape_factor(geo_triangle_file,geo_graph)
    print 'Geo shape ', avg_geo_shape/geo_avg_dist

    del geo_graph

    soc_graph, soc_avg_dist = load_graph(social_trace_file)
    print "average social link length", soc_avg_dist
    avg_soc_shape,x_soc,c_soc = user_shape_factor(social_triangle_file,soc_graph)

    print 'Social shape ', avg_soc_shape/soc_avg_dist

    del soc_graph
    print ''

    plt.figure()
    plt.clf()
    plt.axes(FIG_AXES2)
    plt.semilogx(x,c,'k-')
    plt.semilogx(x_geo,c_geo,'k--')
    plt.semilogx(x_soc,c_soc,'k:')
    plt.legend(['Original data',
        'Geo model','Social model'],loc='upper left')
    #plt.xlabel('Average triangle link length')#size / average link length')
    plt.xlabel(r'$\langle l_{\Delta} \rangle$ [km]')
    plt.ylabel('CDF')
    plt.grid(True)
    #plt.axis((1e-4,10,0,1))
    plt.savefig('%s_null_tri_stats.pdf'%dataset)
    plt.close()
    #sys.exit()
